import torch
from torch import nn
import torch.nn.functional as f

"""
自定义多层感知器模型
nn.Module 集成
init:初始化所有层
forward: 定义模型的运算过程（向前传播过程
"""


class Model(nn.Module):
    """
    HR模型
    """

    def __init__(self):
        super().__init__()
        # 初始化所有层
        self.liner_1 = nn.Linear(20, 64)
        self.liner_2 = nn.Linear(64, 64)
        #  输出层
        self.liner_3 = nn.Linear(64, 1)

    def forward(self, input_value):
        # 第一层 并使用激活函数
        x = f.relu(self.liner_1(input_value))
        #  第二层 并 使用激活函数
        x = f.relu(self.liner_2(x))
        # 输出层 sigmoid输出层激活
        x = f.sigmoid(self.liner_3(x))
        return x

    @staticmethod
    def get_model(learning_rate):
        """
        获取模型
        #  创建自定义模型
        :return: 自定义模型
        """
        model = Model()
        # 对于多层模型，需要使用前序梯度下降法
        opt = torch.optim.Adam(model.parameters(), lr=learning_rate)
        return model, opt
